package com.github.cyf.rateLimit.annotation;

import com.github.cyf.common.exception.BizException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 限流策略，计数器
 * @author chenyifan
 * @create 2025-01-20 17:07
 */
@Aspect
@Component
@Slf4j
public class CounterRateLimitAspect {

    private final ConcurrentHashMap<String, Long> TIMESTAMP = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, AtomicInteger> REQUEST_COUNT = new ConcurrentHashMap<>();

    @Pointcut("@annotation(com.github.cyf.rateLimit.annotation.CounterRateLimit)")
    private void pointCut() {
    }

    @Around("pointCut()")
    public Object interceptor(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method method = methodSignature.getMethod();
        CounterRateLimit counterRateLimit = method.getAnnotation(CounterRateLimit.class);

        int maxRequests = counterRateLimit.requests();
        long windowSizeInMillis = TimeUnit.SECONDS.toMillis(counterRateLimit.timeWindow());

        // 初始化计数器和时间戳
        String methodName = method.toString();
        AtomicInteger count = REQUEST_COUNT.computeIfAbsent(methodName, k -> new AtomicInteger(0));
        Long startTime = TIMESTAMP.computeIfAbsent(methodName, k -> System.currentTimeMillis());

        // 获取当前时间
        long currentTimeMillis = System.currentTimeMillis();
        // 检查时间窗口是否过期
        if (currentTimeMillis - startTime > windowSizeInMillis) {
            // 时间窗口过期，重置计数器和时间戳
            count.set(0);
            TIMESTAMP.put(methodName, currentTimeMillis);
        }
        // 检查请求数量是否超过限制
        if (count.incrementAndGet() > maxRequests) {
            count.decrementAndGet();
            throw new BizException("请求过于频繁，请稍后重试");
        }
        return joinPoint.proceed();
    }
}